from torch import nn


class Tudui(nn.Module):
    def __init__(self):
        super(Tudui, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, stride=1, padding=2, kernel_size=5)
        self.maxpool1 = nn.MaxPool2d(2)
        self.conv2 = nn.Conv2d(32, 32, stride=2, padding=0, kernel_size=5)
        self.maxpool2 = nn.MaxPool2d(2)
        self.conv3 = nn.Conv2d(32, 64, stride=1, padding=2, kernel_size=5)
        self.maxpool3 = nn.MaxPool2d(2)
        self.flatten = nn.Flatten()

        self.linear1 = nn.Linear(1024, 64)
        self.linear2 = nn.Linear(64, 10)
